//#include <boost/asio/ip/udp.hpp>
#include <boost/asio/ip/tcp.hpp>
#include <string>
#include <iostream>
#include <iomanip>
#include <sstream>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>

void dns_test();

void on_send(const boost::system::error_code& error, std::size_t bytes_transferred);
void on_recv(const boost::system::error_code& error, std::size_t bytes_transferred);

class big_uint16_t
{
public:
    big_uint16_t() { }
    big_uint16_t(const big_uint16_t &_src):big_(_src.big_) { }
    big_uint16_t(::uint16_t _little):big_(htons(_little)) { }
    big_uint16_t& operator=(big_uint16_t _src) { big_ = _src.big_; return *this; }
    big_uint16_t& operator=(::uint16_t _little) { big_ = htons(_little); return *this; }
    ::uint16_t value() const { return ntohs(big_); }
    operator ::uint16_t() const { return this->value(); }
private:
    ::uint16_t big_;
};

class big_uint32_t
{
public:
    big_uint32_t() { }
    big_uint32_t(const big_uint32_t &_src):big_(_src.big_) { }
    big_uint32_t(::uint32_t _little):big_(htonl(_little)) { }
    big_uint32_t& operator=(big_uint32_t _src) { big_ = _src.big_; return *this; }
    big_uint32_t& operator=(::uint32_t _little) { big_ = htonl(_little); return *this; }
    ::uint32_t value() const { return ntohl(big_); }
    operator ::uint32_t() const { return this->value(); }
private:
    ::uint32_t big_;
};

namespace dns
{

#pragma pack(push,1)
typedef struct
{
    big_uint16_t id;
    ::uint16_t rd:1;
    ::uint16_t tc:1;
    ::uint16_t aa:1;
    ::uint16_t opcode:4;
    ::uint16_t qr:1;
    ::uint16_t rcode:4;
    ::uint16_t cd:1;
    ::uint16_t ad:1;
    ::uint16_t unused:1;
    ::uint16_t ra:1;
    big_uint16_t qdcount;
    big_uint16_t ancount;
    big_uint16_t nscount;
    big_uint16_t arcount;
} Header;

typedef struct
{
    Header header;
    char mem[512 - sizeof(Header)];
} Packet;

typedef struct
{
    big_uint16_t size;
    Packet packet;
} TcpPacket;

typedef struct
{
    big_uint16_t qtype;
    big_uint16_t qclass;
} QData;

typedef struct _AData : QData
{
    std::string to_string(const Packet *packet) const;
    big_uint32_t ttl;
    big_uint16_t rdlength;
    char rdata[0];
} AData;

typedef struct
{
    std::string to_string(const Packet *packet) const;
    big_uint16_t preference;
    char exchange[0];
} MXData;

typedef struct
{
    big_uint32_t serial;
    big_uint32_t refresh;
    big_uint32_t retry;
    big_uint32_t expire;
    big_uint32_t minimum;
} SoaData;
#pragma pack(pop)

enum QType
{
    DNS_QTYPE_A     =   1,
    DNS_QTYPE_NS    =   2,
    DNS_QTYPE_MD    =   3,
    DNS_QTYPE_MF    =   4,
    DNS_QTYPE_CNAME =   5,
    DNS_QTYPE_SOA   =   6,
    DNS_QTYPE_MB    =   7,
    DNS_QTYPE_MG    =   8,
    DNS_QTYPE_MR    =   9,
    DNS_QTYPE_NULL  =  10,
    DNS_QTYPE_WKS   =  11,
    DNS_QTYPE_PTR   =  12,
    DNS_QTYPE_HINFO =  13,
    DNS_QTYPE_MINFO =  14,
    DNS_QTYPE_MX    =  15,
    DNS_QTYPE_TXT   =  16,
    DNS_QTYPE_RP    =  17,
    DNS_QTYPE_AFSDB =  18,
    DNS_QTYPE_SIG   =  24,
    DNS_QTYPE_KEY   =  25,
    DNS_QTYPE_AAAA  =  28,
    DNS_QTYPE_LOC   =  29,
    DNS_QTYPE_DS    =  43,
    DNS_QTYPE_RRSIG =  46,
    DNS_QTYPE_DNSKEY = 48,
    DNS_QTYPE_AXFR  = 252,
    DNS_QTYPE_MAILB = 253,
    DNS_QTYPE_MAILA = 254,
    DNS_QTYPE_ANY   = 255,
};

enum QClass
{
    DNS_QCLASS_IN   =   1,
    DNS_QCLASS_CS   =   2,
    DNS_QCLASS_CH   =   3,
    DNS_QCLASS_HS   =   4,
    DNS_QCLASS_ANY  = 255,
};

std::string dots_separated_to_query_name(const std::string &name)
{
    std::string res;
    if (name.empty()) {
        return res;
    }
    const char *w = name.c_str();
    while (true) {
        const char *e = w;
        while ((*e != '.') && (*e != '\0')) {
            ++e;
        }
        res += char(e - w);
        res += std::string(w, e - w);
        if (*e == '\0') {
            return res;
        }
        w = e + 1;
    }
}

std::string make_query(const std::string &name, QType qtype, QClass qclass = DNS_QCLASS_IN)
{
    QData qdata;
    qdata.qtype = qtype;
    qdata.qclass = qclass;
    return dots_separated_to_query_name(name) + std::string("", 1) + std::string(reinterpret_cast< const char* >(&qdata), sizeof(qdata));
}

} // namespace dns

std::ostream& operator<<(std::ostream &_out, dns::QType qtype);
std::ostream& operator<<(std::ostream &_out, dns::QClass qclass);
std::ostream& operator<<(std::ostream &_out, const dns::QData &_qdata);
std::ostream& operator<<(std::ostream &_out, const dns::SoaData &_soadata);
std::ostream& operator<<(std::ostream &_out, const dns::Header &_header);
std::ostream& operator<<(std::ostream &_out, const dns::Packet &_packet);

#define DNS_SERVER_ADDR "127.0.0.1"
//#define DNS_SERVER_ADDR "8.8.8.8"
//#define DNS_SERVER_ADDR "172.20.21.254"
//#define DNS_SERVER_ADDR "172.20.21.253"
#define DNS_SERVER_PORT 53
//#define HOST_NAME "iruska.mistr.801.mistrnet.cz"
//#define HOST_NAME "www.google.cz"
#define HOST_NAME "mail.mistrnet.cz"

void dns_test()
{
    boost::asio::io_service io_service;
//    boost::asio::ip::udp::socket s(io_service);
//    boost::asio::ip::udp::endpoint to(boost::asio::ip::address::from_string(DNS_SERVER_ADDR), DNS_SERVER_PORT);
    boost::asio::ip::tcp::socket s(io_service);
    boost::asio::ip::tcp::endpoint to(boost::asio::ip::address::from_string(DNS_SERVER_ADDR), DNS_SERVER_PORT);
    s.open(to.protocol());
    s.connect(to);
    dns::TcpPacket tcp_packet;
    dns::Packet &packet = tcp_packet.packet;
    std::memset(&packet, 0, sizeof(dns::Header));
    packet.header.id = 12345;
    packet.header.rd = 1;
    packet.header.qdcount = 1;
    std::string query = dns::make_query(HOST_NAME, dns::DNS_QTYPE_CNAME);
    std::memcpy(packet.mem, query.c_str(), query.length());
    tcp_packet.size = sizeof(dns::Header) + query.length();
//    s.async_send_to(boost::asio::buffer(&packet, sizeof(dns::Header) + query.length()), to, 0, on_send);
//    s.async_send(boost::asio::buffer(&packet, sizeof(dns::Header) + query.length()), 0, on_send);
    std::cout << s.send(boost::asio::buffer(&tcp_packet, tcp_packet.size + sizeof(tcp_packet.size)), 0) << " bytes sended" << std::endl;
    std::cout << packet << std::endl;
//    boost::asio::ip::udp::endpoint from;
    char data_buffer[0x4000];
    boost::asio::mutable_buffers_1 data(data_buffer, sizeof(data_buffer));
//    s.async_receive_from(data, from, 0, on_recv);
//    s.async_receive(data, 0, on_recv);
    try {
        const ::size_t bytes = s.receive(data, 0);
        std::cout << bytes << " bytes received" << std::endl;
//        s.async_receive(data, 0, on_recv);
//        io_service.run();
        std::cout << reinterpret_cast< dns::TcpPacket* >(data_buffer)->packet << std::endl;
    }
    catch (const std::exception &e) {
        std::cerr << "exception: " << e.what() << std::endl;
        return;
    }
}

void on_send(const boost::system::error_code& error, std::size_t bytes_transferred)
{
    if (!error) {
        std::cout << bytes_transferred << " bytes sended" << std::endl;
    }
    else {
        std::cout << boost::system::system_error(error).what() << std::endl;
    }
}

void on_recv(const boost::system::error_code& error, std::size_t bytes_transferred)
{
    if (!error) {
        std::cout << bytes_transferred << " bytes received" << std::endl;
    }
    else {
        std::cout << boost::system::system_error(error).what() << std::endl;
    }
}

std::string dns_opcode_out(::uint16_t _opcode)
{
    switch (_opcode) {
    case 0:
        return "QUERY ";
    case 1:
        return "IQUERY";
    case 2:
        return "STATUS";
    }
    std::ostringstream o;
    o << std::setw(6) << _opcode;
    return o.str();
}

std::string dns_rcode_out(::uint16_t _rcode)
{
    switch (_rcode) {
    case 0:
        return "OK";
    case 1:
        return "Format error";
    case 2:
        return "Server failure";
    case 3:
        return "Name Error";
    case 4:
        return "Not Implemented";
    case 5:
        return "Refused";
    }
    std::ostringstream o;
    o << _rcode;
    return o.str();
}

std::ostream& operator<<(std::ostream &_out, const dns::Header &_header)
{
    std::ostringstream o;
    o << "ID:      " << std::setw(5) << _header.id << std::endl
      << "FLAGS:   "
      << (_header.qr ? "!" : "?")
      << (_header.rd ? " RD" : " rd")
      << (_header.tc ? " TC" : " tc")
      << (_header.qr ? _header.aa ? " AA" : " aa" : " ..")
      << (_header.cd ? " CD" : " cd")
      << (_header.ad ? " AD" : " ad")
      << (_header.ra ? " RA" : " ra")
      << " " << dns_opcode_out(_header.opcode)
      << " " << dns_rcode_out(_header.rcode);
    return _out << o.str();
}

namespace
{

std::string& dns_name_out(std::string &_out, const char *_base, const char *&_pQ)
{
    bool delimiter = false;
    while (true) {
        const ::uint8_t bytes = *_pQ;
        if (bytes == 0) {
            ++_pQ;
            return _out;
        }
        if (0300 <= bytes) {
            const char *pQ = _base + reinterpret_cast< const big_uint16_t* >(_pQ)->value() - 0xc000;
            _pQ += 2;
            if (delimiter) {
                _out += '.';
            }
            return dns_name_out(_out, _base, pQ);
        }
        if (delimiter) {
            _out += '.';
        }
        else {
            delimiter = true;
        }
        _out += std::string(_pQ + 1, bytes);
        _pQ += 1 + bytes;
    }
}

std::ostream& dns_name_out(std::ostream &_out, const char *_base, const char *&_pQ)
{
    std::string name;
    dns_name_out(name, _base, _pQ);
    name.reserve(18);
    while (name.length() < 18) {
        name += ' ';
    }
    return _out << name;
}

std::string dns_name_out(const char *_base, const char *_pQ)
{
    std::string name;
    dns_name_out(name, _base, _pQ);
    return name;
}

std::string dns_ttl_out(::uint32_t ttl)
{
    std::ostringstream o;
    o << std::setw(2) << std::setfill(' ') << (ttl / 3600) << ':'
      << std::setw(2) << std::setfill('0') << ((ttl % 3600) / 60) << ':'
      << std::setw(2) << std::setfill('0') << (ttl % 60);
    return o.str();
}

std::string dns_data_out(const char *_data, ::size_t _data_len)
{
    std::string o;
    const char *b = _data;
    const char *const end = _data + _data_len;
    while (b < end) {
        const char *c = b;
        while ((c < end) && (0x20 <= *c) && (*c < 0x7f)) {
            ++c;
        }
        o += std::string(b, c - b);
        b = c;
        if (b < end) {
            o += '.';
            ++b;
        }
    }
    return o;
}

}

std::ostream& operator<<(std::ostream &_out, const dns::Packet &_packet)
{
    std::ostringstream o;
    o << _packet.header << std::endl;
    const char *pQ = _packet.mem;
    for (::uint16_t idx = 0; idx < _packet.header.qdcount; ++idx) {
        o << (idx == 0 ? "QUERY:  " : "        ");
        dns_name_out(o, reinterpret_cast< const char* >(&_packet), pQ) << " ";
        o << *reinterpret_cast< const dns::QData* >(pQ) << std::endl;
        pQ += sizeof(dns::QData);
    }
    for (::uint16_t idx = 0; idx < _packet.header.ancount; ++idx) {
        o << (idx == 0 ? "ANSWER: " : "        ");
        dns_name_out(o, reinterpret_cast< const char* >(&_packet), pQ) << " ";
        o << reinterpret_cast< const dns::AData* >(pQ)->to_string(&_packet) << std::endl;
        pQ += sizeof(dns::AData) + reinterpret_cast< const dns::AData* >(pQ)->rdlength.value();
    }
    for (::uint16_t idx = 0; idx < _packet.header.nscount; ++idx) {
        o << (idx == 0 ? "NS:     " : "        ");
        dns_name_out(o, reinterpret_cast< const char* >(&_packet), pQ) << " ";
        o << reinterpret_cast< const dns::AData* >(pQ)->to_string(&_packet) << std::endl;
        pQ += sizeof(dns::AData) + reinterpret_cast< const dns::AData* >(pQ)->rdlength.value();
    }
    for (::uint16_t idx = 0; idx < _packet.header.arcount; ++idx) {
        o << (idx == 0 ? "ADD:    " : "        ");
        dns_name_out(o, reinterpret_cast< const char* >(&_packet), pQ) << " ";
        o << reinterpret_cast< const dns::AData* >(pQ)->to_string(&_packet) << std::endl;
        pQ += sizeof(dns::AData) + reinterpret_cast< const dns::AData* >(pQ)->rdlength.value();
    }
    return _out << o.str();
}

std::ostream& operator<<(std::ostream &_out, dns::QType qtype)
{
    switch (qtype) {
    case dns::DNS_QTYPE_A:
        return _out << "A";
    case dns::DNS_QTYPE_AAAA:
        return _out << "AAAA";
    case dns::DNS_QTYPE_AFSDB:
        return _out << "AFSDB";
    case dns::DNS_QTYPE_ANY:
        return _out << "ANY";
    case dns::DNS_QTYPE_AXFR:
        return _out << "AXFR";
    case dns::DNS_QTYPE_CNAME:
        return _out << "CNAME";
    case dns::DNS_QTYPE_DS:
        return _out << "DS";
    case dns::DNS_QTYPE_DNSKEY:
        return _out << "DNSKEY";
    case dns::DNS_QTYPE_HINFO:
        return _out << "HINFO";
    case dns::DNS_QTYPE_KEY:
        return _out << "KEY";
    case dns::DNS_QTYPE_LOC:
        return _out << "LOC";
    case dns::DNS_QTYPE_MAILA:
        return _out << "MAILA";
    case dns::DNS_QTYPE_MAILB:
        return _out << "MAILB";
    case dns::DNS_QTYPE_MB:
        return _out << "MB";
    case dns::DNS_QTYPE_MD:
        return _out << "MD";
    case dns::DNS_QTYPE_MF:
        return _out << "MF";
    case dns::DNS_QTYPE_MG:
        return _out << "MG";
    case dns::DNS_QTYPE_MINFO:
        return _out << "MINFO";
    case dns::DNS_QTYPE_MR:
        return _out << "MR";
    case dns::DNS_QTYPE_MX:
        return _out << "MX";
    case dns::DNS_QTYPE_NS:
        return _out << "NS";
    case dns::DNS_QTYPE_NULL:
        return _out << "NULL";
    case dns::DNS_QTYPE_PTR:
        return _out << "PTR";
    case dns::DNS_QTYPE_RP:
        return _out << "RP";
    case dns::DNS_QTYPE_RRSIG:
        return _out << "RRSIG";
    case dns::DNS_QTYPE_SIG:
        return _out << "SIG";
    case dns::DNS_QTYPE_SOA:
        return _out << "SOA";
    case dns::DNS_QTYPE_TXT:
        return _out << "TXT";
    case dns::DNS_QTYPE_WKS:
        return _out << "WKS";
    }
    std::ostringstream o;
    o << "\"" << ::uint16_t(qtype) << "\"";
    return _out << o.str();
}

std::ostream& operator<<(std::ostream &_out, dns::QClass qclass)
{
    switch (qclass) {
    case dns::DNS_QCLASS_ANY:
        return _out << "ANY";
    case dns::DNS_QCLASS_CH:
        return _out << "CH";
    case dns::DNS_QCLASS_CS:
        return _out << "CS";
    case dns::DNS_QCLASS_HS:
        return _out << "HS";
    case dns::DNS_QCLASS_IN:
        return _out << "IN";
    }
    return _out << "\"" << ::uint16_t(qclass) << "\"";
}

std::ostream& operator<<(std::ostream &_out, const dns::QData &_qdata)
{
    std::ostringstream o;
    o << std::setw(6) << std::left << dns::QType(_qdata.qtype.value())
      << std::setw(4) << std::left << dns::QClass(_qdata.qclass.value());
    return _out << o.str();
}

std::ostream& operator<<(std::ostream &_out, const dns::SoaData &_soadata)
{
    std::ostringstream o;
    o << std::setw(10) << _soadata.serial.value() << " "
      << dns_ttl_out(_soadata.refresh.value()) << " "
      << dns_ttl_out(_soadata.retry.value()) << " "
      << dns_ttl_out(_soadata.expire.value()) << " "
      << dns_ttl_out(_soadata.minimum.value());
    return _out << o.str();
}

std::string dns::AData::to_string(const dns::Packet *packet) const
{
    std::ostringstream o;
    o << *static_cast< const dns::QData* >(this) << " " << dns_ttl_out(this->ttl.value());
    switch (this->qtype.value()) {
    case dns::DNS_QTYPE_A:
        if (this->rdlength.value() == sizeof(struct ::in_addr)) {
            o << " " << ::inet_ntoa(*reinterpret_cast< const struct ::in_addr* >(this->rdata));
            return o.str();
        }
        o << " " << this->rdlength.value() << "bytes";
        return o.str();
    case dns::DNS_QTYPE_AAAA:
        if (this->rdlength.value() == sizeof(struct ::in6_addr)) {
            char buf[INET6_ADDRSTRLEN + 1];
            o << " " << ::inet_ntop(AF_INET6, this->rdata, buf, sizeof(buf));
            return o.str();
        }
        o << " " << this->rdlength.value() << "bytes";
        return o.str();
    case dns::DNS_QTYPE_NS:
    case dns::DNS_QTYPE_CNAME:
        o << " " << dns_name_out(reinterpret_cast< const char* >(packet), this->rdata);
        return o.str();
    case dns::DNS_QTYPE_MX:
        o << " " << reinterpret_cast< const dns::MXData* >(this->rdata)->to_string(packet);
        return o.str();
    case dns::DNS_QTYPE_SOA:
        const char *pQ = this->rdata;
        dns_name_out(o << " ", reinterpret_cast< const char* >(packet), pQ);
        dns_name_out(o << " ", reinterpret_cast< const char* >(packet), pQ);
        o << " " << *reinterpret_cast< const dns::SoaData* >(pQ);
        return o.str();
    }
    o << " " << this->rdlength.value() << "bytes " << dns_data_out(this->rdata, this->rdlength);
    return o.str();
}

std::string dns::MXData::to_string(const dns::Packet *packet) const
{
    std::ostringstream o;
    o << std::setw(3) << this->preference.value() << " " << dns_name_out(reinterpret_cast< const char* >(packet), this->exchange);
    return o.str();
}
